#!/usr/bin/perl
$dir = "/data4/bsi/RandD/Workflow/genotyping_1000genome_impute";

#PARAMETERS TO BE CHANGED
my $dirtemp = "/data4/bsi/RandD/Workflow/temp";
$rounded = "mariza_sample";
my $impute_ref = "/data4/bsi/RandD/Workflow/genotyping_1000genome_impute/impute_ref1";
#$impute_ref = "/data4/bsi/refdata/genetics/impute/ALL_1000G_phase1integrated_feb2012_impute/";
my $impute_window = 5000000;
my $impute_edge = 125;#in KB
chomp($impute_ref);
chomp($impute_window);
chomp($impute_edge);
#checking impute window and edge sizes
if(!($impute_window >= 2000000 && $impute_window <= 5000000))
{
	die "impute window should be between 2 and 5 mb\n";
}
if($impute_edge != 125)
{
	die "impute edge should be 125 in kb\n";
}
$dirtemp =~ s/\/$//g;
$impute_ref =~ s/\/$//g;

#looking for directories from shapeit
undef(@check_chr);
@ls = `ls -d $dirtemp/$rounded/*`;
for($i=0;$i<@ls;$i++)
{
        chomp($ls[$i]);
        @array = split(/\//,$ls[$i]);
        $chr = pop(@array);
        if($chr =~ m/^\d+$/ && $chr <24)
        {
				#print "$chr\n";
                push(@check_chr,$chr);
        }
		undef(@array);
        #print "$ls[$i]\n";
}
@check_chr = sort {$a <=> $b} @check_chr;
#die "test @check_chr\n";
#running impute on these files steps
#creating the xchr imputation sample file sex information
if($check_chr[@check_chr-1] == 23)
{
	print "sucess\n";
	open(XCHR,"$dirtemp/$rounded/23/snps_chr23.fam")or die "no file exists $dirtemp/$rounded/23/snps_chr23.fam\n";
	while(<XCHR>)
	{
		chomp($_);
		@array = split(" ",$_);
		$xchr{$array[1]} = $array[4];
	}
	close(XCHR);
	open(XCHR,"$dirtemp/$rounded/23/snps_chr23.sample")or die "no file exists $dirtemp/$rounded/23/snps_chr23.sample\n";
	open(WRXCHR,">$dirtemp/$rounded/23/snps_chr23_1.sample")or die "unable to write file $dirtemp/$rounded/23/snps_chr23_1.sample\n";
	$line = <XCHR>;
	chomp($line);
	print WRXCHR $line." sex\n";
	$line = <XCHR>;
	chomp($line);
	print WRXCHR $line." D\n";
	while(<XCHR>)
	{
			chomp($_);
		@array = split(" ",$_);
		if(!exists($xchr{$array[1]}))
		{
			die "plink file doesn't match with shapeit program output sample file\n";
		}
		print WRXCHR "$_ $xchr{$array[1]}\n";
	}
	system("mv $dirtemp/$rounded/23/snps_chr23_1.sample $dirtemp/$rounded/23/snps_chr23.sample");
	undef %xchr;
	close(XCHR);
	close(WRXCHR);
}

#creating the array job file for the impute
open(ARRAY_IMPUTE,">$dirtemp/$rounded/ArrayJob_file_impute") or die "unable to create file $dirtemp/$rounded/ArrayJob_file_impute\n";
#calculating window sizes for all chrs
for($i=0;$i<@check_chr;$i++)
#for($i=0;$i<1;$i++)
{
	print "\n\n\n\n\nDEALING WITH CHR $check_chr[$i]\n\n\n\n";
	open(BIM,"$dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].haps") or die "$dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].haps file not exists\n";
	$line = <BIM>;
	@array = split(" ",$line);
	$start = $array[2];
	while($line = <BIM>)
	{
		$line1 = $line;
	}
	@array = split(" ",$line1);
	$stop = $array[2];
	$num_windows = ($stop-$start+1)/$impute_window;
	$rem = ($stop-$start+1)%$impute_window;
	$num_windows =~ s/\.\d+//g;
	if($rem != 0)
	{
        	$num_windows =$num_windows+1;
	}
	print "start: $start stop $stop numwindows :  $num_windows\n";
	#creating windows start and stop
	for($j=0;$j<$num_windows;$j++)
	{
		$count_snp_window_indata[$j] = 0;
		$count_snp_window_reference[$j] = 0;
		if($j ==0)
		{
			$start_window[$j] = $start;
			$stop_window[$j] = $start_window[$j]+ $impute_window -1;
		}
		else
		{
			$start_window[$j] = $stop_window[$j-1]+1;
			$stop_window[$j] = $start_window[$j]+ $impute_window -1;
		}
	}
	#print "@start_window\n@stop_window\n";
	#counting the number of input data markers in each window
	open(BIM,"$dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].haps") or die "$dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].haps file not exists\n";
	$j = 0;
	while($line = <BIM>)
	{
		@array = split(" ",$line);
		if($array[2] >= $start_window[$j] && $array[2] <= $stop_window[$j])
		{
			$count_snp_window_indata[$j]++;
		}	
		else
		{
			$j++;
			if($array[2] >= $start_window[$j] && $array[2] <= $stop_window[$j])
			{
				$count_snp_window_indata[$j]++;
			}
		}
	}
	#print "$num\n";
	#counting the number of reference data marers in each window
	if($check_chr[$i] ==23)
	{
		open(BIM,"gunzip -c $impute_ref/ALL_1000G_phase1integrated_feb2012_chrX_nonPAR_impute.legend.gz|") or die "no file found $impute_ref/ALL_1000G_phase1integrated_feb2012_chrX_nonPAR_impute.legend.gz\n";
	}
	else
	{
		$file = "ALL_1000G_phase1integrated_feb2012_chr$check_chr[$i]".'_impute.legend.gz';;
		open(BIM,"gunzip -c $impute_ref/$file|") or die "no file found $impute_ref/$file\n";
	}
	$line = <BIM>;
	$j = 0;
	#print $line."\n";
	$k = 0;
	while($line = <BIM>)
	{
		@array = split(" ",$line);
		#start: storing starting value of the chr in the reference
		if($k == 0)
		{
			$start_chr_ref = $array[1];
		}
		$stop_chr_ref = $array[1];
		$k++;
		#end: storing starting value of the chr in the reference
		if($array[1] < $start_window[0] || $array[1] > $stop_window[@stop_window -1])
		{

		}	
		elsif($array[1] <= $stop_window[$j] && $array[1] >= $start_window[$j])
		{
			$count_snp_window_reference[$j]++;
		}
		else
		{
			$j++;
			if($array[1] >= $start_window[$j] && $array[1] <= $stop_window[$j])
                        {
				$count_snp_window_reference[$j]++;
			}
		}
	}
	print "\n\n raw counts of windows\n";
	for($j=0;$j<@count_snp_window_indata;$j++)
        {
                print "Start : $start_window[$j] Stop : $stop_window[$j] Data : $count_snp_window_indata[$j] Reference : $count_snp_window_reference[$j]\n";
	}
	#craeting zeros for windows count zero
	for($j=0;$j<@count_snp_window_indata;$j++)
	{
		#print "Start : $start_window[$j] Stop : $stop_window[$j] Data : $count_snp_window_indata[$j] Reference : $count_snp_window_reference[$j]\n";
		#checking for minimum 200 count snps
		if($count_snp_window_indata[$j] < 200 ||  $count_snp_window_reference[$j] < 200)
		{
				$stop_window[$j]  = 0;
				$start_window[$j+1] = 0;
				if($j == @count_snp_window_indata -1)
				{
					$start_window[$j] = 0;
					$stop_window[$j-1]  = 0;
				}
		}
	}
	print "\n\n\n\n\n\n";
	print "\n\n temp counts of windows\n";
	for($j=0;$j<@count_snp_window_indata;$j++)
        {
                print "Start : $start_window[$j] Stop : $stop_window[$j] Data : $count_snp_window_indata[$j] Reference : $count_snp_window_reference[$j]\n";
        }

	$k = -1;
	#$start_window_new[0] = $start_window[0];
	#merging windows based on zeros
	for($j=0;$j<@count_snp_window_indata;$j++)
        {
		if($start_window[$j] != 0 && $stop_window[$j] != 0)
		{
			$k++;
			$start_window_new[$k] = $start_window[$j];
			$stop_window_new[$k] = $stop_window[$j];
			$count_snp_window_indata_new[$k] = $count_snp_window_indata[$j];
			$count_snp_window_reference_new[$k] = $count_snp_window_reference[$j];
		}
		elsif($start_window[$j] ==0)
		{
                        $stop_window_new[$k] = $stop_window[$j];
                        $count_snp_window_indata_new[$k] =$count_snp_window_indata_new[$k]+$count_snp_window_indata[$j];
                        $count_snp_window_reference_new[$k] =$count_snp_window_reference_new[$k]+$count_snp_window_reference[$j];
                }
		else
		{
			$k++;
			$start_window_new[$k] = $start_window[$j];
			$count_snp_window_indata_new[$k] =$count_snp_window_indata[$j];
			$count_snp_window_reference_new[$k] =$count_snp_window_reference[$j]; 
		} 
	}
	#writing parameters in to array job impute file
	#start:extending windows start and stop according to reference
	$start_window_new[0]= $start_chr_ref;
	$stop_window_new[@stop_window_new -1] = $stop_chr_ref; 
	#start:extending windows start and stop according to reference
	print "\n\n\n\n\n\n";
	print "\n finalized windows and counts\n";
	for($j=0;$j<=$k;$j++)
        {
                print "Start : $start_window_new[$j] Stop : $stop_window_new[$j] Data : $count_snp_window_indata_new[$j] Reference : $count_snp_window_reference_new[$j]\n";
				#if window is greater than 5 mb attach -allow_large_regions
				$length_window = $stop_window_new[$j] - $start_window_new[$j] +$impute_edge*2*1000;
				print "$length_window\n";
				if($length_window > 5000000)
				{
					$option_val = '-use_prephased_g -allow_large_regions';
				}
				else
				{
					$option_val = '-use_prephased_g ';
				}
        	if($check_chr[$i] != 23)
		{
			$map = "genetic_map_chr$check_chr[$i]"."_combined_b37.txt";
			$legend = "ALL_1000G_phase1integrated_feb2012_chr$check_chr[$i]"."_impute.legend.gz";
			$hap = "ALL_1000G_phase1integrated_feb2012_chr$check_chr[$i]"."_impute.hap.gz";
			print ARRAY_IMPUTE "$option_val -m $impute_ref/$map -h $impute_ref/$hap -l $impute_ref/$legend -int $start_window_new[$j] $stop_window_new[$j] -Ne 20000  -buffer $impute_edge -known_haps_g $dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].haps -o $dirtemp/$rounded/$check_chr[$i]/impute_out_chr$check_chr[$i].window$j\n";
		}
		else
		{
			$map = "genetic_map_chrX_nonPAR_combined_b37.txt";
                        $legend = "ALL_1000G_phase1integrated_feb2012_chrX_nonPAR_impute.legend.gz";
                        $hap = "ALL_1000G_phase1integrated_feb2012_chrX_nonPAR_impute.hap.gz";
						#$option_val =~ s/-use_prephased_g//g;
			print ARRAY_IMPUTE "$option_val -chrX -m $impute_ref/$map -h $impute_ref/$hap -l $impute_ref/$legend -int $start_window_new[$j] $stop_window_new[$j] -Ne 20000  -buffer $impute_edge -known_haps_g $dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].haps -sample_known_haps_g $dirtemp/$rounded/$check_chr[$i]/snps_chr$check_chr[$i].sample -o $dirtemp/$rounded/$check_chr[$i]/impute_out_chr$check_chr[$i].window$j\n";
			#print ARRAY_IMPUTE "$option_val -chrX -m $impute_ref/$map -h $impute_ref/$hap -l $impute_ref/$legend -int $start_window_new[$j] $stop_window_new[$j] -Ne 20000  -buffer $impute_edge -g $dirtemp/$rounded/$check_chr[$i]/snps_chr$i.genotype -sample_g $dirtemp/$rounded/$check_chr[$i]/snps_chr$i.sample -o $dirtemp/$rounded/$check_chr[$i]/impute_out_chr$check_chr[$i].window$j\n";
			
		}
	}
	undef(@start_window);
	undef(@stop_window);
	undef(@start_window_new);
	undef(@stop_window_new);
	undef(@array);
	undef(@count_snp_window_indata_new);
	undef(@count_snp_window_indata);
	undef(@count_snp_window_reference_new);
	undef(@count_snp_window_reference);
}
#submit the impute jobs
open(ARRAY_IMPUTE,"$dir/bin/ArrayJob_impute.csh") or die "no file exists ArrayJob_impute.csh\n";
@impute = <ARRAY_IMPUTE>;
$impute =join("",@impute);
@impute =();
@impute = split('zzz',$impute);
system("mkdir $dirtemp/$rounded/impute_logfiles_sungrid");
$count_impute_jobs = `wc -l $dirtemp/$rounded/ArrayJob_file_impute`;
$count_impute_jobs =~  s/ .*//g;
chomp($count_impute_jobs);
$impute =$impute[0].$count_impute_jobs.$impute[1]."$dirtemp/$rounded/impute_logfiles_sungrid".$impute[2]."$dirtemp/$rounded/impute_logfiles_sungrid".$impute[3]."$dirtemp/$rounded/ArrayJob_file_impute".$impute[4]."$dir/bin/impute_v2.2.2_x86_64_static/impute2 ".$impute[5];
open(ARRAY_SHAPEIT,">$dirtemp/$rounded/ArrayJob_impute.csh") or die "unable to create the array job file\n";
print ARRAY_SHAPEIT $impute;
close(ARRAY_SHAPEIT);
close(ARRAY_IMPUTE);
#submitting and storing the job id
system("qsub $dirtemp/$rounded/ArrayJob_impute.csh &> $dirtemp/$rounded/jobid_impute");
##readin job id from submit_shapeit
open(ARRAY_IMPUTE,"$dirtemp/$rounded/jobid_impute") or die "unable to open file $dirtemp/$rounded/jobid_impute\n";
$impute = <ARRAY_IMPUTE>;
print "$impute\n";
@impute =split(" ",$impute);
@impute1 =split(/\./,$impute[2]);
print "JOB ID $impute1[0]\n";
$job_id_imputeit = $impute1[0];

#making the array job to wait
open(ARRAY_SHAPEIT,">$dirtemp/$rounded/ArrayJob_shapeit_wait.csh") or die "unable to create the array job wait file shape it \n";
print ARRAY_SHAPEIT '#! /bin/bash'."\n";
print ARRAY_SHAPEIT '#$ -q 1-day'."\n";
print ARRAY_SHAPEIT '#$ -l h_vmem=5G'."\n";
print ARRAY_SHAPEIT '#$ -M prodduturi.naresh@mayo.edu'."\n";
print ARRAY_SHAPEIT '#$ -m a'."\n";
print ARRAY_SHAPEIT '#$ -V'."\n";
print ARRAY_SHAPEIT '#$ -cwd'."\n";
$com = '#$';
print ARRAY_SHAPEIT "$com -e $dirtemp/$rounded/\n";
print ARRAY_SHAPEIT "$com -o $dirtemp/$rounded/\n";
if(-e "$dirtemp/$rounded/waiting.txt")
{
	system("rm $dirtemp/$rounded/waiting.txt");
}
print ARRAY_SHAPEIT "cp $dirtemp/$rounded/ArrayJob_impute.csh $dirtemp/$rounded/waiting.txt\n";

$sys = "qsub -hold_jid $job_id_imputeit $dirtemp/$rounded/ArrayJob_shapeit_wait.csh\n";
print $sys."\n";
system($sys);

$flag = 1;
while($flag > 0)
{
        if($flag%100000000 == 0)
        {
                $flag = 1;
                print "waiting for the shapeit array job to complete\n";
        }
        $flag++;
        if(-e "$dirtemp/$rounded/waiting.txt")
        {
                $flag =0;
        }
}
if(-e "$dirtemp/$rounded/waiting.txt")
{
        system("rm $dirtemp/$rounded/waiting.txt");
}
#checking all the impute jobs are successfully completed
open(ARRAY_IMPUTE,"$dirtemp/$rounded/ArrayJob_file_impute") or die "unable to find file $dirtemp/$rounded/ArrayJob_file_impute\n";
open(IMPUTE_3_PROB,">$dirtemp/$rounded/Combined_impute_results_3_prob_allchr")  or die "unable to create file $dirtemp/$rounded/Combined_impute_results_3_prob_allchr\n";
$num_samp =0;
while(<ARRAY_IMPUTE>)
{
	chomp($_);
	$_ =~ s/\s+$//g;
	undef(@array);
	undef(@array1);
	@array = split(" ",$_);
	@array1 = split(/\//,$array[@array-1]);
	$chr = $array1[@array1-2];
	open(BUFF,"$array[@array-1]")or die "file not found after imputatio $_\n";
	while($line = <BUFF>)
	{
		if($num_samp ==0)
		{
			@array2 = split(" ",$line);
			$num_samp = @array2;
			$header = "chr SNPID marker POS alleleA alleleB";
			$num_samp =($num_samp-5)/3;
			for($i=1;$i<=$num_samp;$i++)
			{
				$header=$header." $i $i $i";
			}
			print IMPUTE_3_PROB $header."\n";
		
		}

		$line = $chr." $line";
				#die $line."\n";
		print IMPUTE_3_PROB $line; 
	}
	#system("rm $array[@array-1] ");
}
close(IMPUTE_3_PROB);
open(WRBUFF,">$dirtemp/$rounded/beagle_r2") or die "not able to write $dirtemp/$rounded/beagle_r2\n";
print WRBUFF "marker_identifier\tminor allele\tminor allele frequency\tallelic r-squared\tdosage r-squared\tHWE dosage r-squared\taccuracy\tmissing score\n";
close(WRBUFF);
system("cut -d ' ' -f3,5- $dirtemp/$rounded/Combined_impute_results_3_prob_allchr|java -jar $dir/bin/gprobsmetrics.jar >> $dirtemp/$rounded/beagle_r2");


